import torch
import pandas as pd
# from maskrcnn_benchmark.structures import BoxList

# 加载预测结果
preds = torch.load('/home/ai999/project/div-align-dg/runs/thyroid_r101_c4_voc/inference0/thyroid_test/predictions.pth')

# 初始化一个列表，用于存储所有的预测框信息
results = []

# 遍历每张图像的预测结果
for idx, pred in enumerate(preds):
    boxes = pred.bbox  # 预测框的位置 (xyxy格式)
    scores = pred.get_field('scores')  # 预测框的分数
    labels = pred.get_field('labels')  # 预测框的类别标签

    # 将每个预测框的信息存储到results中
    for i in range(len(boxes)):
        result = {
            'image_id': idx,
            'box': boxes[i].cpu().numpy(),
            'score': scores[i].cpu().numpy(),
            'label': labels[i].cpu().numpy()
        }
        results.append(result)

# 将结果转换为DataFrame
df = pd.DataFrame(results)

# 导出到Excel
df.to_excel('thyroid_r101_c4_voc-predictions_results.xlsx', index=False)

print(f'Prediction results have been saved to "predictions_results.xlsx".')
